"""LLM Plays Pokemon 提示词管理模块"""

import os
import json
import logging
from typing import Dict, Any, List, Optional

from llm_pokemon.knowledge.base import KnowledgeBase

logger = logging.getLogger(__name__)


class PromptManager:
    """提示词管理类，负责组织和管理与LLM的交互提示"""
    
    def __init__(self, max_turns: int, history_path: str):
        """初始化提示词管理器
        
        Args:
            max_turns: 触发总结的最大对话轮数
            history_path: 会话历史存储路径
        """
        self.max_turns = max_turns
        self.history_path = history_path
        self.history = []
        self.turn_count = 0
        
        # 加载历史记录
        self._load_history()
        
        logger.info("提示词管理器初始化完成")
    
    def _load_history(self):
        """加载会话历史"""
        if os.path.exists(self.history_path):
            try:
                with open(self.history_path, 'r', encoding='utf-8') as f:
                    self.history = json.load(f)
                self.turn_count = len([msg for msg in self.history if msg.get("role") == "assistant"])
                logger.info(f"已加载会话历史: {self.history_path}, 当前对话轮数: {self.turn_count}")
            except Exception as e:
                logger.error(f"加载会话历史失败: {e}")
                self.history = []
                self.turn_count = 0
        else:
            logger.info(f"会话历史文件不存在: {self.history_path}，创建新会话")
            self.history = []
            self.turn_count = 0
    
    def save_history(self):
        """保存会话历史"""
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(os.path.abspath(self.history_path)), exist_ok=True)
            
            # 保存历史记录
            with open(self.history_path, 'w', encoding='utf-8') as f:
                json.dump(self.history, f, indent=2, ensure_ascii=False)
            
            logger.info(f"已保存会话历史: {self.history_path}")
        except Exception as e:
            logger.error(f"保存会话历史失败: {e}")
    
    def add_user_message(self, content: str):
        """添加用户消息
        
        Args:
            content: 消息内容
        """
        self.history.append({
            "role": "user",
            "content": content
        })
    
    def add_assistant_message(self, content: str):
        """添加助手消息
        
        Args:
            content: 消息内容
        """
        self.history.append({
            "role": "assistant",
            "content": content
        })
        self.turn_count += 1
    
    def add_system_message(self, content: str):
        """添加系统消息
        
        Args:
            content: 消息内容
        """
        self.history.append({
            "role": "system",
            "content": content
        })
    
    def clear_history(self):
        """清空会话历史"""
        self.history = []
        self.turn_count = 0
        logger.info("已清空会话历史")
    
    def get_history(self) -> List[Dict[str, str]]:
        """获取会话历史
        
        Returns:
            会话历史列表
        """
        return self.history.copy()
    
    def needs_summarization(self) -> bool:
        """检查是否需要进行总结
        
        Returns:
            是否需要总结
        """
        return self.turn_count >= self.max_turns
    
    def compose_prompt(self, knowledge_base: KnowledgeBase) -> str:
        """组成提示词
        
        Args:
            knowledge_base: 知识库实例
            
        Returns:
            完整提示词
        """
        # 系统提示词
        system_prompt = self._get_system_prompt()
        
        # 知识库内容
        knowledge_content = self._get_knowledge_content(knowledge_base)
        
        # 工具定义
        tools_definition = self._get_tools_definition()
        
        # 会话历史
        history_content = self._format_history()
        
        # 组合完整提示词
        prompt = f"{system_prompt}\n\n{knowledge_content}\n\n{tools_definition}\n\n{history_content}"
        
        return prompt
    
    def compose_summarization_prompt(self, history: List[Dict[str, str]], knowledge_base: KnowledgeBase) -> str:
        """组成总结提示词
        
        Args:
            history: 会话历史
            knowledge_base: 知识库实例
            
        Returns:
            总结提示词
        """
        # 系统总结提示
        system_prompt = """请总结以下对话内容，重点关注：
1. 游戏进度和重要事件
2. 玩家当前状态（位置、队伍、物品等）
3. 已完成的任务和未完成的目标
4. 重要的游戏信息和发现

总结应简洁明了，但包含所有关键信息。"""
        
        # 知识库内容
        knowledge_content = self._get_knowledge_content(knowledge_base)
        
        # 格式化历史记录
        history_content = "\n\n".join([f"{msg['role']}: {msg['content']}" for msg in history])
        
        # 组合总结提示词
        prompt = f"{system_prompt}\n\n{knowledge_content}\n\n对话历史：\n{history_content}"
        
        return prompt
    
    def _get_system_prompt(self) -> str:
        """获取系统提示词
        
        Returns:
            系统提示词
        """
        return """你是一个Pokemon游戏助手，能够通过模拟器与Pokemon游戏交互。你可以使用以下工具：

1. use_emulator - 发送按钮指令到游戏模拟器
2. navigator - 在游戏中导航到指定位置或执行特定操作
3. update_knowledge_base - 更新你的长期记忆

注意事项：
- 不要完全依赖视觉，经常使用知识库获取信息
- 保持对游戏状态的跟踪，包括位置、队伍和进度
- 提供清晰的游戏建议和策略
- 当遇到新的Pokemon、NPC或地点时，更新知识库"""
    
    def _get_knowledge_content(self, knowledge_base: KnowledgeBase) -> str:
        """获取知识库内容
        
        Args:
            knowledge_base: 知识库实例
            
        Returns:
            格式化的知识库内容
        """
        sections = knowledge_base.list_sections()
        content_parts = []
        
        for section_id in sections:
            section_content = knowledge_base.get_section(section_id)
            if section_content:
                content_parts.append(f"<section id=\"{section_id}\">\n{section_content}\n</section>")
        
        if content_parts:
            return "知识库内容：\n" + "\n\n".join(content_parts)
        else:
            return "知识库为空"
    
    def _get_tools_definition(self) -> str:
        """获取工具定义
        
        Returns:
            工具定义字符串
        """
        return """可用工具：

1. use_emulator(buttons: List[str], pause_ms: Optional[int] = None) -> Dict
   - buttons: 按钮序列，如["A", "B", "Up", "Down", "Left", "Right", "Start", "Select"]
   - pause_ms: 执行后暂停的毫秒数（可选）
   - 返回：包含截图、叠加层和游戏状态的结果

2. navigator(target: Dict) -> Dict
   - target: 目标信息，可以是以下之一：
     - {"coordinates": [x, y]} - 导航到指定坐标
     - {"scene": "scene_name"} - 导航到指定场景
     - {"action": "action_name"} - 执行特定操作
   - 返回：导航结果和状态信息

3. update_knowledge_base(operation: str, section_id: str, content: Optional[str] = None) -> Dict
   - operation: 操作类型，可选值：'add', 'update', 'delete', 'get'
   - section_id: 部分ID
   - content: 部分内容，仅在add和update操作时需要
   - 返回：操作结果"""
    
    def _format_history(self) -> str:
        """格式化会话历史
        
        Returns:
            格式化的会话历史
        """
        if not self.history:
            return "对话历史为空"
        
        formatted_messages = []
        for msg in self.history:
            role = msg["role"]
            content = msg["content"]
            
            if role == "user":
                formatted_messages.append(f"用户: {content}")
            elif role == "assistant":
                formatted_messages.append(f"助手: {content}")
            elif role == "system":
                formatted_messages.append(f"系统: {content}")
        
        return "对话历史：\n" + "\n\n".join(formatted_messages)